Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[TIR] Check producer predicate in ReverseComputeInline #13338

Merged
merged 5 commits into from
Nov 10, 2022
Merged

[TIR] Check producer predicate in ReverseComputeInline #13338

merged 5 commits into from
Nov 10, 2022

Conversation

masahi
Copy link
Member

@masahi masahi commented Nov 9, 2022

Currently, even if a producer block has a non-trivial predicate, ReverseComputeInline goes ahead and produces buggy a schedule. The real-world example I've hit was the following. After ReverseComputeInline the compute_4 block, the predicate ((ax1_0 * 4 + ax1_1) * 32 + ax1_2) * 2 + ax1_3 < 64) disappears. See https://gist.github.com/masahi/01a80b86062122ad57b9b1fd785fb960 for a repro.

         ...
         for ax0, ax1_0, ax1_1, ax1_2, ax1_3 in T.grid(32, 1, 4, 32, 2):
             with T.block("conv2d_nhwc_reindex_shared"):
                 T.where(((ax1_0 * 4 + ax1_1) * 32 + ax1_2) * 2 + ax1_3 < 64)
                 v0 = T.axis.spatial(50176, ax2_0_0_ax3_0_0_fused // 4 * 6272 + ax2_0_1_ax3_0_1_fused * 32 + ax0)
                 v1 = T.axis.spatial(256, ax2_0_0_ax3_0_0_fused % 4 * 64 + (ax1_0 * 256 + ax1_1 * 64 + ax1_2 * 2 + ax1_3))
                 T.reads(p7[()], conv2d_nhwc_reindex_shared[v0, v1], p2[0, 0, 0, v1], p3[0, 0, 0, v1], p4[v1], p5[v1], p6[v1], p8[0])
                 T.writes(compute_3[v0 // 3136, v0 % 3136 // 56, v0 % 56, v1])
                 compute_3[v0 // 3136, v0 % 3136 // 56, v0 % 56, v1] = T.q_multiply_shift(T.max(T.min(p7[()] + T.q_multiply_shift_per_axis(conv2d_nhwc_reindex_shared[v0, v1] - p2[0, 0, 0, v1] + p3[0, 0, 0, v1], p4[v1], p5[v1], p6[v1], 31, False, True, dtype="int32"), 255), 0) - p8[0], 1457846997, 31, 0, dtype="int32")
 for i0_12, i1_12, i2_12, i3_12 in T.grid(16, 56, 56, 256):
     with T.block("compute_4"):
         i0_13, i1_13, i2_13, i3_13 = T.axis.remap("SSSS", [i0_12, i1_12, i2_12, i3_12])
         T.reads(compute_3[i0_13, i1_13, i2_13, i3_13], p9[i0_13, i1_13, i2_13, i3_13])
         T.writes(compute[i0_13, i1_13, i2_13, i3_13])
         compute[i0_13, i1_13, i2_13, i3_13] = T.max(T.min(compute_3[i0_13, i1_13, i2_13, i3_13] + T.q_multiply_shift(p9[i0_13, i1_13, i2_13, i3_13], 2101000910, 31, 0, dtype="int32"), 255), 0)

So we should disallow ReverseComputeInine when the producer has a non-trivial predicate. But if the predicate in the new inlined block can imply the original predicate in the producer block, we can still allow ReverseComputeInline to be applied. These two cases are demonstrated in the test cases.

@vinx13 @junrushao @Hzfengsy

@tvm-bot
Copy link
Collaborator

tvm-bot commented Nov 9, 2022 •
edited
Loading

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

Generated by tvm-bot

@masahi masahi force-pushed the tir-reverse-inline-check-predicate branch from 8de97bc to da4bc25 Compare November 10, 2022 00:02
junrushao
junrushao approved these changes Nov 10, 2022
View reviewed changes
Copy link
Member

@junrushao junrushao left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That makes a lot of sense to me :-) Thanks for the PR!

@masahi masahi merged commit 6d9d213 into apache:main Nov 10, 2022
xinetzone pushed a commit to daobook/tvm that referenced this pull request Nov 10, 2022
@masahi
[TIR] Check producer predicate in ReverseComputeInline (apache#13338)
3097867 
* [TIR] Disallow reverse inline into a producer with non-trivial predicate

* add test

* Allow cases where the producer predicate can be implied by the new
predicate of the inlined block

* remove unused variable

* update comment in test to reflect the change in ReverseComputeInline
xinetzone pushed a commit to daobook/tvm that referenced this pull request Nov 25, 2022
@masahi @xinetzone
[TIR] Check producer predicate in ReverseComputeInline (apache#13338)
e319476 
* [TIR] Disallow reverse inline into a producer with non-trivial predicate

* add test

* Allow cases where the producer predicate can be implied by the new
predicate of the inlined block

* remove unused variable

* update comment in test to reflect the change in ReverseComputeInline
@leandron leandron mentioned this pull request Feb 1, 2023
Closed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@vinx13 vinx13 vinx13 approved these changes

@junrushao junrushao junrushao approved these changes

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@masahi @tvm-bot @vinx13 @junrushao